import sys
import typing
import os

from utils.MainWindow import *
import utils.FetchDialog as FetchDialog
import utils.SubDIalog as SubDIalog
import utils.EnhanceVDIalog as EnhanceVDIalog
import utils.OutputDialog as OutputDialog
import utils.DistributeDialog as DistributeDialog
import utils.RepeatfillDialog as RepeatfillDialog

from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QTransform
from PyQt5.QtCore import Qt, QPoint
from PyQt5 import QtCore
from PyQt5.QtWidgets import QApplication , QMainWindow, QGraphicsScene, QGraphicsView
from PyQt5.QtWidgets import QFileDialog, QLabel, QListWidget, QListWidgetItem

from utils.tools import *

import re

def contains_chinese(path):
    """检测路径是否含中文"""
    pattern = re.compile(r'[\u4e00-\u9fa5]')
    return pattern.search(path) is not None

def ndarray_to_QPixmap(ndarray):
    height, width, channel = ndarray.shape
    bytesPerLine = 3 * width
    qImg = QImage(ndarray.data, width, height, bytesPerLine, QImage.Format_BGR888).rgbSwapped()
    pixmap = QPixmap(qImg)
    return pixmap

class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)

        self.datasets = []



        #debug
        tmp = dataset("utils/image.jpg")
        self.widget.start_plot(tmp)

        #打开文件夹
        self.actionopendir.triggered.connect(self.open_folder_dialog)
        #往当前数据集添加图片
        self.actionadd.triggered.connect(self.add_image)
        #往当前数据集添加文件夹
        self.actionadddir.triggered.connect(self.add_dir)

        #选择指定的数据集文件夹
        self.DatasetfileListWidget.itemClicked.connect(self.show_imagelist)        
        #选择指定的图片
        self.ImagefileListWidget.itemClicked.connect(self.show_image)
        #选择指定的stage
        self.StageListWidget.itemClicked.connect(self.show_stage_image)        
        #点击上下选择图片
        self.ImagefileListWidget.itemSelectionChanged.connect(self.on_item_selected)
        self.StageListWidget.itemSelectionChanged.connect(self.on_item_selected)

        #提取操作
        self.FetchDialog = FetchDialog.FetchDialog()
        self.FetchPushButton.clicked.connect(self.FetchDialog.exec_)
        self.FetchDialog.fetch_signal.connect(self.Fetch)        

        #保存或者输出操作
        self.OutputDialog = OutputDialog.OutputDialog()
        self.actionsave.triggered.connect(self.OutputDialog.exec_)
        self.OutputPushButton.clicked.connect(self.OutputDialog.exec_)
        self.OutputDialog.output_signal.connect(self.Output)
        

        #分发操作
        self.DistributeDialog = DistributeDialog.DistributeDialog()
        self.DistributePushButton.clicked.connect(self.DistributeDialog.exec_)
        self.DistributeDialog.output_signal.connect(self.Distribute)

        #repeat_fill操作
        self.repeat_fillDialog = RepeatfillDialog.RepeatfillDialog()
        self.RepeatFillPushButton.clicked.connect(self.repeat_fillDialog.exec_)
        self.repeat_fillDialog.output_signal.connect(self.repeat_fill)

        #去重操作
        self.UniquePushButton.clicked.connect(self.Unique)

        #减去数据集操作
        self.SubPushButton.clicked.connect(self.show_Sub_Dialog)

        #验证关键点操作
        self.VerityKptsPushButton.clicked.connect(self.VerifyKpts)

        #计算TV操作
        self.CalcTVPushButton.clicked.connect(self.CalcTV)

        #生成HDT操作
        self.G.clicked.connect(self.GenerateHDT)

        #水平翻转操作
        self.openDirectoryPushButton_14.clicked.connect(self.Fliplr)

        #亮度增强操作
        self.EnhanceVDialog = EnhanceVDIalog.EnhanceVDialog()
        self.openDirectoryPushButton_15.clicked.connect(self.EnhanceVDialog.exec_)
        self.EnhanceVDialog.output_signal.connect(self.EnhanceV)
    
    def open_folder_dialog(self):
        """打开文件夹对话框"""
        folder_path = QFileDialog.getExistingDirectory()
        if folder_path:
            self.DatasetfileListWidget.addItem(folder_path)
            da = Datasets(folder_path)
            self.datasets.append(da)

            #将当前ImagefileListWidget的索引设置为最后一个
            self.ImagefileListWidget.setCurrentRow(
                self.ImagefileListWidget.count() - 1)
            self.show_imagelist()

            #检查当前路径是否包含中文
            if contains_chinese(folder_path):
                QtWidgets.QMessageBox.warning(self, "提示", "路径包含中文,可能会导致有些功能无法使用", QtWidgets.QMessageBox.Yes)


    def add_image(self):
        """往当前数据集再加载图片"""
        image_path, _ = QFileDialog.getOpenFileName(None, "选择图片", "", "Images (*.png *.jpg)")
        
        if not image_path:
            return
        elif contains_chinese(image_path):
            QtWidgets.QMessageBox.warning(self, "提示", "图片路径错误,请检查是否包含中文", QtWidgets.QMessageBox.Yes)
        else:
            image = cv2.imread(image_path)
            da = Datasets(image)
            da.data[0].img_path = image_path
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.load(da)

            self.show_imagelist()
            self.ImagefileListWidget.setCurrentRow(self.ImagefileListWidget.count() - 1)
            self.show_image(self.ImagefileListWidget.currentItem())
            
            QtWidgets.QMessageBox.information(self, "提示", "添加图片完成", QtWidgets.QMessageBox.Yes)
            
    def add_dir(self):
        """添加文件夹"""
        folder_path = QFileDialog.getExistingDirectory()
        if folder_path:
            self.DatasetfileListWidget.addItem(folder_path)
            da = Datasets(folder_path)
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.load(da)

            self.show_imagelist()
            self.ImagefileListWidget.setCurrentRow(self.ImagefileListWidget.count() - 1)
            self.show_image(self.ImagefileListWidget.currentItem())

            if contains_chinese(folder_path):
                QtWidgets.QMessageBox.warning(self, "提示", "路径包含中文,可能会导致有些功能无法使用", QtWidgets.QMessageBox.Yes)
            else:
                QtWidgets.QMessageBox.information(self, "提示", "添加文件夹完成", QtWidgets.QMessageBox.Yes)

    def on_item_selected(self):
        """点击上下显示图片"""
        list_widget = self.sender()
        if not isinstance(list_widget, QListWidget):
            return
        self.show_image(list_widget.currentItem())
        
    def show_stage(self, current_dataset):
        """显示stage中的内容"""
        self.StageListWidget.clear()
        for i in current_dataset.stage:
            stage_imgpath = current_dataset.data[i].img_path
            self.StageListWidget.addItem(str(stage_imgpath))

    def show_imagelist(self):
        """选择指定的数据集时显示图片在列表中"""
        current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]

        #选择指定的数据集时显示图片在列表中
        self.ImagefileListWidget.clear()
        for file in current_dataset.data:
            self.ImagefileListWidget.addItem(str(file.img_path))

        #显示stage中的内容
        self.show_stage(current_dataset)           

        #显示数据集的构成        
        _translate = QtCore.QCoreApplication.translate
        self.labelingbox.setTitle(_translate("MainWindow", "数据集构成"))
        self.labelListWidget.clear()
        composition = current_dataset.statistics()
        for infomat in composition:
            self.labelListWidget.addItem(infomat)

    def show_image(self, item):
        """将选择的图片展示在中间"""
        
        #获取当前item的索引值
        if(self.DatasetfileListWidget.currentRow() > len(self.datasets) - 1):
            return
        index = self.ImagefileListWidget.currentRow()
        current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
        data = current_dataset.data[index]

        self.widget.start_plot(data)

    def show_stage_image(self, item):
        """将选择的stage展示在中间"""

        for child in self.widget.children():
            if isinstance(child, QLabel):
                child.deleteLater()
        
        self.scale = 1.0
        #获取当前item的索引值
        if(self.DatasetfileListWidget.currentRow() > len(self.datasets) - 1):
            return
        index = self.StageListWidget.currentRow()
        current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
        data = current_dataset.data[index]

        self.widget.start_plot(data)

    def Output(self, OutputConfig):
        """输出数据集"""
        if len(OutputConfig):
            if(len(self.datasets)):
                current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
                current_dataset.output(
                OutputConfig['save_path'], OutputConfig['label_format'], OutputConfig['val_rate'], 
                OutputConfig['add'], OutputConfig['reserve_name'], OutputConfig['move'], 
                OutputConfig['sum'], OutputConfig['pop'], OutputConfig['show_log'], 
                OutputConfig['toDatasets'], OutputConfig['generate_hdt'], OutputConfig['same_dir'], 
                OutputConfig['reset_name'], OutputConfig['data_name'])

    def Fetch(self, FetchConfig):
        """提取指定数据集到stage"""
        if(FetchConfig['type'] != 'dont_fetch'):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]

            if(FetchConfig['type'] == 'class'):
                current_dataset.fetch_data_core(
                    FetchConfig['color'], FetchConfig['ts'],
                    FetchConfig['class'], FetchConfig['sum'])
            
            elif(FetchConfig['type'] == 'dis'):
                current_dataset.fetch_data_dis(
                    FetchConfig['thred'], FetchConfig['net'])
                
            elif(FetchConfig['type'] == 'V'):
                current_dataset.fetch_data_V(
                    FetchConfig['thred'])
                
            elif(FetchConfig['type'] == 'source'):
                current_dataset.fetch_data_source(
                    FetchConfig['source'])
            
            self.show_stage(current_dataset)

    def Distribute(self, DistributeConfig):
        """分发数据集"""
        if(len(DistributeConfig)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.distribute(
                DistributeConfig['path'], DistributeConfig['sum_datasets'], DistributeConfig['sum_all'], 
                DistributeConfig['sum_each'], DistributeConfig['add'], DistributeConfig['reserve_name'], 
                DistributeConfig['move'], DistributeConfig['toDatasets'], DistributeConfig['same_dir']) 
            #弹出一个对话框，显示分发完成
            QtWidgets.QMessageBox.information(self, "提示", "分发完成", QtWidgets.QMessageBox.Yes)

    def repeat_fill(self, repeat_fillConfig):
        """重复填充"""
        if(len(repeat_fillConfig)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.repeat_fill(
                repeat_fillConfig['low_limit'], repeat_fillConfig['unique'], 
                repeat_fillConfig['class_limit'])
            
            self.show_imagelist()
            self.show_stage(current_dataset)
            
            #弹出一个对话框，显示重复填充完成            
            QtWidgets.QMessageBox.information(self, "提示", "重复填充完成", QtWidgets.QMessageBox.Yes)

    def Unique(self):
        """去重"""
        if(len(self.datasets)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.unique()

            self.show_imagelist()
            self.show_stage(current_dataset)

    def show_Sub_Dialog(self):
        """显示减去数据集对话框"""
        if(len(self.datasets)):
            self.SubDialog = SubDIalog.SubDialog()
            self.SubDialog.DatasetList.clear()

            #将数据集列表中的数据集添加到对话框中
            for i in range(self.DatasetfileListWidget.count()):
                self.SubDialog.DatasetList.addItem(self.DatasetfileListWidget.item(i).text())
            
            self.SubDialog.show()
            self.SubDialog.output_signal.connect(self.Sub)

    def Sub(self, SubConfig):
        """减去指定数据集"""
        if(len(SubConfig)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            sub_dataset = Datasets(SubConfig)
            current_dataset.sub(sub_dataset)

            self.show_imagelist()
            self.show_stage(current_dataset)

            #弹出一个对话框，显示减去完成            
            QtWidgets.QMessageBox.information(self, "提示", "减去完成", QtWidgets.QMessageBox.Yes)

    def VerifyKpts(self):
        if(len(self.datasets)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            informat = current_dataset.verify_kps()

            self.show_imagelist()
            self.show_stage(current_dataset)

            data_sum = informat['data_sum']
            error_sum = informat['error_sum']
            QtWidgets.QMessageBox.information(self, "提示", 
                        f'Got {data_sum} targets in total, {error_sum} targets out of bounds, {data_sum - error_sum} targets in confirm.'
                        , QtWidgets.QMessageBox.Yes)

    def CalcTV(self):
         """计算TV值"""
         if(len(self.datasets)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.build_TV()

            self.show_imagelist()
            self.show_stage(current_dataset)
            QtWidgets.QMessageBox.information(self,
                 "提示", "计算完成", QtWidgets.QMessageBox.Yes)

    def GenerateHDT(self):
        """生成HDT""" 
        if(len(self.datasets)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.generate_hdt()

            self.show_imagelist()
            self.show_stage(current_dataset)
            QtWidgets.QMessageBox.information(self,
             "提示", "生成完成", QtWidgets.QMessageBox.Yes)

    def Fliplr(self):
        """水平翻转"""
        if(len(self.datasets)):
            current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
            current_dataset.fliplr()

            self.show_imagelist()
            self.show_stage(current_dataset)
            QtWidgets.QMessageBox.information(self,
             "提示", "翻转完成", QtWidgets.QMessageBox.Yes)
    
    def EnhanceV(self, rate):
        if(rate > 0):
            if(len(self.datasets)):
                current_dataset = self.datasets[self.DatasetfileListWidget.currentRow()]
                current_dataset.enhance_V(rate)

                self.show_imagelist()
                self.show_stage(current_dataset)
                QtWidgets.QMessageBox.information(self,
                "提示", "增强完成", QtWidgets.QMessageBox.Yes)        

if __name__ == '__main__':
    app = QApplication(sys.argv)
    myWin = MainWindow()
    myWin.show()
    sys.exit(app.exec_())